################################################################################
class rap_states_ca(rap_states):
  """ Store set of configuration-averaged states. """
  def __init__(self,filename=None,**args):
    self.clear()
    if filename: self.loadfile(filename,**args)

# ------------------------------------------------------------------------------

  def clear(self):
    """ Erase all data. """
    rap_states.clear(self)
    self.v_cup='CA'
    self.v_conf={}

# ------------------------------------------------------------------------------

  def clear_levels(self):
    """ Just erase level data. """
    rap_states.clear_levels(self)
    self.v_cup='CA'
    self.v_conf={}

# ------------------------------------------------------------------------------

  def loadfile(self,filename,**args):
    """ Load in level information from level data file. """
    ft=rlf_filetype(filename)
    if ft == 'x006':
      self.load_xstar(filename,**args)
    else:
      msg='do not know how to read file: %s'%filename
      raise ErrorRAP(msg)

# ------------------------------------------------------------------------------

  def reindex(self,remap=None):
    """ 
    Reassign level indices.  Indices are remapped according to dictionary,
    remap, old (key) -> new (value).  If remap is None, then reindex according
    to energy order.
    """
    imap=rap_states.reindex(self,remap)
    self.v_conf=dict([(imap[i],val) for (i,val) in self.v_conf.items()])

# ------------------------------------------------------------------------------

  def configurations(self,levlist=None,imap=False):
    """ Return list of all configurations. """
    if levlist is None: levlist=self.indices()
    if imap is True: levlist=map(self._mapme,levlist)
    out=[]
    for i in levlist:
      if self.v_conf[i] in out: continue
      out.append(self.v_conf[i])
    return out

# ------------------------------------------------------------------------------

  def configuration(self,idx,imap=False):
    """ Return configuration of level with given index. """
    tidx=self._mapme(idx,imap)
    return self.v_conf[tidx]

# ------------------------------------------------------------------------------

  def level(self,idx,imap=False):
    """ Return all information for single level (conf,par,E) """
    tidx=self._mapme(idx,imap)
    out=(self.v_conf[tidx],self.v_parity[tidx], self.v_energy[tidx])
    return out

# ------------------------------------------------------------------------------

  def label(self,idx,imap=False,spacer=None):
    """ Return level label; optional spacer can be used instead of space. """
    tidx=self._mapme(idx,imap)
    out=str(self.v_conf[tidx])
    if spacer is not None: out=out.replace(' ',spacer)
    return out

# ------------------------------------------------------------------------------

  def weight(self,idx,imap=False):
    """ Return weight of level. """
    tidx=self._mapme(idx,imap)
    return self.v_conf[tidx].weight()

# ------------------------------------------------------------------------------

  def add_level(self,conf,energy,idx=False):
    """ Add new level to list """
    if idx is not False and idx in self:
      msg='index already exists: %s'%idx
      raise ErrorRAP(msg)
    tidx=idx
    if tidx is False: 
      tidx=1
      if len(self) > 0: tidx=max(self.indices())+1
    tconf=conf
    if not isinstance(tconf,rap_config): tconf=rap_config(tconf)
    if self.num_electrons() is None:
      self.v_nelec=tconf.num_electrons()
    else:
      if tconf.num_electrons() != self.num_electrons():
        msg='inconsistent number of electrons (expected: %s); conf: %s'%\
            (self.num_electrons(),tconf)
        raise ErrorRAP(msg)
    if tconf in self.v_conf.values():
      msg='duplicate state found: %s'%tconf
      raise ErrorRAP(msg)
    self.v_conf[tidx]=tconf
    self.v_parity[tidx]=tconf.parity()
    self.v_energy[tidx]=energy

# ------------------------------------------------------------------------------

  def add_levels(self,toadd,idxstart=None):
    """ 
    Add set of levels to structure. 

    The input list, toadd, contains tuples with: energy (Ry), index, and
    configuration.  If the indexing is to be automated, then set index in the
    tuples to -1.
    """
    if len(toadd) == 0: return
    toadd.sort()
    idxcur=1
    if len(self) > 0: idxcur=max(self.indices())+1
    if idxstart is not None:
      if idxstart < idxcur: 
        msg='invalid starting index: %s'%idxstart
        raise ErrorRAP(msg)
      idxcur=idxstart
    tmp={}
    nechk=None
    for (en,idx,conf) in toadd:
      tidx=idx
      if tidx in self or tidx in tmp: 
        msg='duplicate index: %s'%tidx
        raise ErrorRAP(msg)
      if tidx < 0:
        tidx=idxcur
        idxcur+=1
      tconf=conf
      if not isinstance(conf,rap_config): tconf=rap_config(conf)
      if nechk is None: nechk=tconf.num_electrons()
      if tconf.num_electrons() != nechk:
        msg='inconsistent number of electrons (expected: %s); conf: %s'%\
            (nechk,tconf)
        raise ErrorRAP(msg)
      tmp[tidx]=(en,tconf)
    self.v_nelec=nechk
    for idx,(en,conf) in tmp.iteritems():
      if conf in self.v_conf.values():
        msg='duplicate state found: %s'%conf
        raise ErrorRAP(msg)
      self.v_conf[idx]=conf
      self.v_parity[idx]=conf.parity()
      self.v_energy[idx]=en

# ------------------------------------------------------------------------------

  def set_lmap(self,other):
    """ Map to other rap_states_ca instance. """
    if not isinstance(other,rap_states_ca):
      msg='argument must be of type rap_states_ca'
      raise ErrorRAP(msg)
    if self.num_electrons() != other.num_electrons():
      msg='inconsistent number of electrons; expected/have: %s, %s'%\
          (self.num_electrons(),other.num_electrons())
      raise ErrorRAP(msg)
    tmap={}
    for i in other:
      k=None
      for j in self:
        if self[j] == other[i]:
          k=j
          break
      tmap[i]=k
    self.v_lmap=tmap
    for j in other:
      if j not in self.v_lmap: self.v_lmap[j]=None
    self.v_lmap_r=dict([(v,k) for k,v in self.v_lmap.items() if v is not None])
    for i in self:
      if i not in self.v_lmap_r: self.v_lmap_r[i]=None

# ------------------------------------------------------------------------------

  def cfilter(self,clist,subset=None,imap=False):
    """ Filter list of states based on configurations in clist. """
    if not isinstance(clist,list) and not isinstance(clist,tuple): clist=[clist]
    tclist=[]
    for c in clist:
      tc=c
      if not isinstance(tc,rap_config): tc=rap_config(tc)
      tclist.append(tc)

    idx=self.indices()
    if subset is not None:
      idx=subset[:]
      if imap is True: idx=self.lmap(idx)

    out=[]
    for i in idx:
      if self.configuration(i) in tclist: out.append(i)
    if imap is True: return self.lmap_reverse(out)
    return out

# ------------------------------------------------------------------------------

  def load_xstar(self,filename):
    """ 
    Load level information from xstar type 006 file.
    """
    # column id's
    id_energy=7
    id_j2=8
    id_ip=10
    id_g=12
    id_l=13
    id_z=14
    id_ion=16
    id_id=15
    id_conf=17

    # read level list from file
    ion=None
    ne=None
    ldata=[]
    have=[]
    xid0=None
    fin=open(filename,'r')
    for line in fin:
      if line.strip() == '': continue
      data=line.split()

      confstr=data[id_conf-1]
      xid=int(data[id_id-1])
      ion=int(data[id_ion-1])
      g=abs(int(data[id_g-1]))
      l=int(data[id_l-1])
      j2=int(float(data[id_j2-1]))-1
      w=g*(2*l+1)
      if g == 1 or l == 0:
        cup='XX'
      elif w < j2+1:
        cup='CA'
        if 'CA' not in have: have.append('CA')
      elif w == j2+1:
        cup='LS'
        if 'LS' not in have: have.append('LS')
      else:
        j2min=abs(g-1-2*l)
        j2max=g-1+2*l
        okay=False
        for j2t in range(j2min,j2max+1,2):
          if j2t == j2: okay=True
        if okay:
          cup='IC'
          if 'IC' not in have: have.append('IC')
        else:
          cup='??'
      e=float(data[id_energy-1])
      if e == 0.0:
        ne=rap_config(confstr).num_electrons()
        zeff=int(data[id_z-1])
        ion_thresh=float(data[id_ip-1])
      chk=confstr.lower()
      if 'su' in chk or 'sp' in chk: continue 
      if 'co' in chk or 'cn' in chk: continue
      if 'fake' in chk: continue
      ldata.append([e/RYDBERG,cup,confstr,g,l,j2,xid])
    fin.close()
    if ne is None: raise RuntimeError('cannot find ground state')

    # id's for ldata
    id_e=0
    id_cup=1
    id_conf=2
    id_g=3
    id_l=4
    id_j2=5
    id_xid=6

    # get number of electrons from ground state and standardize configurations
    for i in range(len(ldata)):
      ldata[i][id_conf]=rap_config(ldata[i][id_conf],ne=ne)

    # get total weight of each configuration
    weight={}
    wtot={}
    cups={}
    for i in range(len(ldata)):
      conf=str(ldata[i][id_conf])
      if conf not in wtot: wtot[conf]=ldata[i][id_conf].weight()
      if conf not in cups: cups[conf]=[]
      if ldata[i][id_cup] not in cups[conf]: cups[conf].append(ldata[i][id_cup])
      if ldata[i][id_cup] in ['XX','LS']:
        if conf not in weight:
          weight[conf]=ldata[i][id_g]*(2*ldata[i][id_l]+1)
        else:
          weight[conf]+=ldata[i][id_g]*(2*ldata[i][id_l]+1)
      else:
        if conf not in weight:
          weight[conf]=(ldata[i][id_j2]+1)
        else:
          weight[conf]+=(ldata[i][id_j2]+1)

    # check if configurations have a common coupling scheme
    for c in cups.keys():
      tmp=cups[c][:]
      if 'XX' in tmp: tmp.remove('XX')
      if len(tmp) > 1:
        raise RuntimeError('mixed coupling scheme not supported')
      if tmp == [] and len(have) == 1: tmp=have[:]
      cups[c]=tmp[:]
      if cups[c] == []: cups[c]=['CA']

    # check that 'CA' configurations have all possible levels
    for c in weight.keys():
      if 'CA' not in cups[c]: continue
      elif weight[c] > wtot[c]:
        raise RuntimeError('configuration, '+str(c)+', has too many levels')

    # add levels having CA coupling
    toadd=[]
    for lev in ldata:
      idx=lev[id_xid]
      conf=lev[id_conf]
      sc=str(conf)
      if 'CA' not in cups[sc] and cups[sc] != ['XX']: continue
      energy=lev[id_e]
      toadd.append((energy,idx,conf))
    if len(toadd) == 0:
      msg='no CA data found in file, %s'%filename
      raise ErrorRAP(msg)
    self.set_nuclear_charge(zeff)
    self.set_ionization_threshold(ion_thresh/RYDBERG)
    self.add_levels(toadd)

# ------------------------------------------------------------------------------

  def write_xstar(self,filename,subset=None,precision=5):
    """ 
    Write xstar type 6 file. Must have ionization potential and nuclear charge
    specified.
    """
    if self.ionization_threshold() == None:
      raise ErrorRAP('must define ionization threshold')
    if self.nuclear_charge() == None:
      raise ErrorRAP('must define nuclear charge')
    indices=self.indices()
    if subset is not None: indices=subset

    zeff=self.nuclear_charge()
    ne=self.num_electrons()
    ionno=xstar_ionencode(zeff,ne)

    aform='%'+str(8+precision)+'.'+str(precision)+'E'
    iform='%6i'
    prefix=iform%6+iform%13+iform%0+iform%4+iform%6
    fout=open(filename,'w')
    for i in indices:
      conf=self.configuration(i)
      e=self.energy(i)*RYDBERG
      wgt=float(self.weight(i))
      ip=self.ionization_threshold()*RYDBERG
      np=-1.0
      if ip > 0.0: 
        z1=float(zeff-ne+1)
        de=abs(self.ionization_threshold()-self.energy(i))
        np=z1/sqrt(de)
      nval=0
      if conf is not None: nval=rlf_ssdecode(conf.valence_subshell())[0]
      g,l=0,-1
      label=self.label(i,spacer='.')
      llen=len(label)
      line=prefix+iform%llen+' '+aform%e+aform%wgt+aform%np+aform%ip+' '+\
                  iform%nval+iform%g+iform%l+iform%zeff+iform%i+\
                  iform%ionno+' '+label+' %'
      fout.write(line+'\n')
    fout.close()

